#include "stdafx.h"
#include "net/FTPClientSession.hpp"
#include "net/SocketStream.hpp"
#include "net/ServerSocket.hpp"
#include "charset/ascii.hpp"
using JHCPP::charset::CAscii;
#include "utils/numberformatter.hpp"
using JHCPP::utils::CNumberFormatter;

NET_NAMESPACE_BEGIN

CFTPClientSession::CFTPClientSession()
	: m_port(0),
	m_pControlSocket(0),
	m_pDataStream(0),
	m_passiveMode(true),
	m_fileType(TYPE_BINARY),
	m_supports1738(true),
	m_serverReady(false),
	m_isLoggedIn(false),
	m_timeout(DEFAULT_TIMEOUT)
{
}

CFTPClientSession::CFTPClientSession(const CStreamSocket& socket)
	: m_host(socket.address().host().toString()),
	m_port(socket.address().port()),
	m_pControlSocket(new CDialogSocket(socket)),
	m_pDataStream(0),
	m_passiveMode(true),
	m_fileType(TYPE_BINARY),
	m_supports1738(true),
	m_serverReady(false),
	m_isLoggedIn(false),
	m_timeout(DEFAULT_TIMEOUT)
{
	m_pControlSocket->setReceiveTimeout(m_timeout);
}


CFTPClientSession::CFTPClientSession(const std::string& host,
								   UInt16 port,
								   const std::string& username,
								   const std::string& password)
	: m_host(host),
	m_port(port),
	m_pControlSocket(new CDialogSocket(CSocketAddress(host, port))),
	m_pDataStream(0),
	m_passiveMode(true),
	m_fileType(TYPE_BINARY),
	m_supports1738(true),
	m_serverReady(false),
	m_isLoggedIn(false),
	m_timeout(DEFAULT_TIMEOUT)
{
	if (!username.empty())
		login(username, password);
	else
		m_pControlSocket->setReceiveTimeout(m_timeout);
}

CFTPClientSession::~CFTPClientSession()
{
	try
	{
		close();
	}
	catch (...) 
	{
	}
}

void CFTPClientSession::setTimeout(const CTimeSpan& timeout)
{
	if (!isOpen())
		throw FTPException("Connection is closed.");

	m_timeout = timeout;
	m_pControlSocket->setReceiveTimeout(timeout);
}

CTimeSpan CFTPClientSession::getTimeout() const
{
	return m_timeout;
}

void CFTPClientSession::setPassive(bool flag, bool useRFC1738)
{
	m_passiveMode  = flag;
	m_supports1738 = useRFC1738;
}

bool CFTPClientSession::getPassive() const
{
	return m_passiveMode;
}

void CFTPClientSession::open(const std::string& host,
							UInt16 port,
							const std::string& username,
							const std::string& password)
{
	m_host = host;
	m_port = port;
	if (!username.empty())
	{
		login(username, password);
	}
	else
	{
		m_pControlSocket = new CDialogSocket(CSocketAddress(m_host, m_port));
		m_pControlSocket->setReceiveTimeout(m_timeout);
	}
}

void CFTPClientSession::login(const std::string& username, const std::string& password)
{
	if (m_isLoggedIn) logout();

	int status = FTP_POSITIVE_COMPLETION * 100;
	std::string response;
	if (!m_pControlSocket)
	{
		m_pControlSocket = new CDialogSocket(CSocketAddress(m_host, m_port));
		m_pControlSocket->setReceiveTimeout(m_timeout);
	}

	if (!m_serverReady)
	{
		status = m_pControlSocket->receiveStatusMessage(response);
		if (!isPositiveCompletion(status))
			throw FTPException("Cannot login to server", response, status);

		m_serverReady = true;
	}

	status = sendCommand("USER", username, response);
	if (isPositiveIntermediate(status))
		status = sendCommand("PASS", password, response);
	if (!isPositiveCompletion(status)) 
		throw FTPException("Login denied", response, status);

	setFileType(m_fileType);
	m_isLoggedIn = true;
}

void CFTPClientSession::logout()
{
	if (!isOpen())
		throw FTPException("Connection is closed.");

	if (m_isLoggedIn)
	{
		try { endTransfer(); }
		catch (...) { }
		std::string response;
		sendCommand("QUIT", response);
		m_isLoggedIn = false;
	}
}

void CFTPClientSession::close()
{
	logout();
	m_pControlSocket->close();
	delete m_pControlSocket;
	m_pControlSocket = 0;
	m_serverReady = false;
}

void CFTPClientSession::setFileType(CFTPClientSession::FileType type)
{
	std::string response;
	int status = sendCommand("TYPE", (type == TYPE_TEXT ? "A" : "I"), response);
	if (!isPositiveCompletion(status)) throw FTPException("Cannot set file type", response, status);
	m_fileType = type;
}

CFTPClientSession::FileType CFTPClientSession::getFileType() const
{
	return m_fileType;
}

std::string CFTPClientSession::systemType()
{
	std::string response;
	int status = sendCommand("SYST", response);
	if (isPositiveCompletion(status))
		return response.substr(4);
	else
		throw FTPException("Cannot get remote system type", response, status);
}

void CFTPClientSession::setWorkingDirectory(const std::string& path)
{
	std::string response;
	int status = sendCommand("CWD", path, response);
	if (!isPositiveCompletion(status))
		throw FTPException("Cannot change directory", response, status);
}

std::string CFTPClientSession::getWorkingDirectory()
{
	std::string response;
	int status = sendCommand("PWD", response);
	if (isPositiveCompletion(status))
		return extractPath(response);
	else
		throw FTPException("Cannot get current working directory", response, status);
}

void CFTPClientSession::cdup()
{
	std::string response;
	int status = sendCommand("CDUP", response);
	if (!isPositiveCompletion(status))
		throw FTPException("Cannot change directory", response, status);
}

void CFTPClientSession::rename(const std::string& oldName, const std::string& newName)
{
	std::string response;
	int status = sendCommand("RNFR", oldName, response);
	if (!isPositiveIntermediate(status)) 
		throw FTPException(std::string("Cannot rename ") + oldName, response, status);
	status = sendCommand("RNTO", newName, response);
	if (!isPositiveCompletion(status)) 
		throw FTPException(std::string("Cannot rename to ") + newName, response, status);
}

void CFTPClientSession::remove(const std::string& path)
{
	std::string response;
	int status = sendCommand("DELE", path, response);
	if (!isPositiveCompletion(status)) 
		throw FTPException(std::string("Cannot remove " + path), response, status);
}

void CFTPClientSession::createDirectory(const std::string& path)
{
	std::string response;
	int status = sendCommand("MKD", path, response);
	if (!isPositiveCompletion(status)) 
		throw FTPException(std::string("Cannot create directory ") + path, response, status);
}

void CFTPClientSession::removeDirectory(const std::string& path)
{
	std::string response;
	int status = sendCommand("RMD", path, response);
	if (!isPositiveCompletion(status)) 
		throw FTPException(std::string("Cannot remove directory ") + path, response, status);
}

std::istream& CFTPClientSession::beginDownload(const std::string& path)
{
	if (!isOpen())
		throw FTPException("Connection is closed.");

	delete m_pDataStream;
	m_pDataStream = 0;
	m_pDataStream = new CSocketStream(establishDataConnection("RETR", path));
	return *m_pDataStream;
}

void CFTPClientSession::endDownload()
{
	endTransfer();
}

std::ostream& CFTPClientSession::beginUpload(const std::string& path)
{
	if (!isOpen())
		throw FTPException("Connection is closed.");

	delete m_pDataStream;
	m_pDataStream = 0;
	m_pDataStream = new CSocketStream(establishDataConnection("STOR", path));
	return *m_pDataStream;
}

void CFTPClientSession::endUpload()
{
	endTransfer();
}

std::istream& CFTPClientSession::beginList(const std::string& path, bool extended)
{
	if (!isOpen())
		throw FTPException("Connection is closed.");

	delete m_pDataStream;
	m_pDataStream = 0;
	m_pDataStream = new CSocketStream(establishDataConnection(extended ? "LIST" : "NLST", path));
	return *m_pDataStream;
}

void CFTPClientSession::endList()
{
	endTransfer();
}

void CFTPClientSession::abort()
{
	if (!isOpen())
		throw FTPException("Connection is closed.");

	m_pControlSocket->sendByte(CDialogSocket::TELNET_IP);
	m_pControlSocket->synch();
	std::string response;
	int status = sendCommand("ABOR", response);
	if (status == 426)
		status = m_pControlSocket->receiveStatusMessage(response);
	if (status != 226) 
		throw FTPException("Cannot abort transfer", response, status);
}

int CFTPClientSession::sendCommand(const std::string& command, std::string& response)
{
	if (!isOpen())
		throw FTPException("Connection is closed.");

	m_pControlSocket->sendMessage(command);
	return m_pControlSocket->receiveStatusMessage(response);
}

int CFTPClientSession::sendCommand(const std::string& command, const std::string& arg, std::string& response)
{
	if (!isOpen())
		throw FTPException("Connection is closed.");

	m_pControlSocket->sendMessage(command, arg);
	return m_pControlSocket->receiveStatusMessage(response);
}

std::string CFTPClientSession::extractPath(const std::string& response)
{
	std::string path;
	std::string::const_iterator it  = response.begin();
	std::string::const_iterator end = response.end();
	while (it != end && *it != '"') ++it;
	if (it != end)
	{
		++it;
		while (it != end)
		{
			if (*it == '"')
			{
				++it;
				if (it == end || (it != end && *it != '"')) break;
			}
			path += *it++;
		}
	}
	return path;
}

CStreamSocket CFTPClientSession::establishDataConnection(const std::string& command, const std::string& arg)
{
	if (m_passiveMode)
		return passiveDataConnection(command, arg);
	else
		return activeDataConnection(command, arg);
}

CStreamSocket CFTPClientSession::activeDataConnection(const std::string& command, const std::string& arg)
{
	if (!isOpen())
		throw FTPException("Connection is closed.");

	CServerSocket server(CSocketAddress(m_pControlSocket->address().host(), 0));
	sendPortCommand(server.address());
	std::string response;
	int status = sendCommand(command, arg, response);
	if (!isPositivePreliminary(status)) 
		throw FTPException(command + " command failed", response, status);
	if (server.poll(m_timeout, CSocket::SELECT_READ))
		return server.acceptConnection();
	else
		throw FTPException("The server has not initiated a data connection");
}

CStreamSocket CFTPClientSession::passiveDataConnection(const std::string& command, const std::string& arg)
{
	CSocketAddress sa(sendPassiveCommand());
	CStreamSocket sock(sa);
	std::string response;
	int status = sendCommand(command, arg, response);
	if (!isPositivePreliminary(status)) 
		throw FTPException(command + " command failed", response, status);
	return sock;
}

void CFTPClientSession::sendPortCommand(const CSocketAddress& addr)
{
	if (m_supports1738)
	{
		if (sendEPRT(addr))
			return;
		else
			m_supports1738 = false;
	}
	sendPORT(addr);
}

CSocketAddress CFTPClientSession::sendPassiveCommand()
{
	CSocketAddress addr;
	if (m_supports1738)
	{
		if (sendEPSV(addr))
			return addr;
		else
			m_supports1738 = false;
	}
	sendPASV(addr);
	return addr;
}

bool CFTPClientSession::sendEPRT(const CSocketAddress& addr)
{
	std::string arg("|");
	arg += addr.af() == AF_INET ? '1' : '2';
	arg += '|';
	arg += addr.host().toString();
	arg += '|';
	arg += CNumberFormatter::format(addr.port());
	arg += '|';
	std::string response;
	int status = sendCommand("EPRT", arg, response);
	if (isPositiveCompletion(status))
		return true;
	else if (isPermanentNegative(status))
		return false;
	else
		throw FTPException("EPRT command failed", response, status);
}

void CFTPClientSession::sendPORT(const CSocketAddress& addr)
{
	std::string arg(addr.host().toString());
	for (std::string::iterator it = arg.begin(); it != arg.end(); ++it)
	{
		if (*it == '.') *it = ',';
	}
	arg += ',';
	UInt16 port = addr.port();
	arg += CNumberFormatter::format(port/256);
	arg += ',';
	arg += CNumberFormatter::format(port % 256);
	std::string response;
	int status = sendCommand("PORT", arg, response);
	if (!isPositiveCompletion(status)) 
		throw FTPException("PORT command failed", response, status);
}

bool CFTPClientSession::sendEPSV(CSocketAddress& addr)
{
	std::string response;
	int status = sendCommand("EPSV", response);
	if (isPositiveCompletion(status))
	{
		parseExtAddress(response, addr);
		return true;
	}
	else if (isPermanentNegative(status))
	{
		return false;
	}
	else throw FTPException("EPSV command failed", response, status);
}

void CFTPClientSession::sendPASV(CSocketAddress& addr)
{
	std::string response;
	int status = sendCommand("PASV", response);
	if (!isPositiveCompletion(status)) 
		throw FTPException("PASV command failed", response, status);
	parseAddress(response, addr);
}

void CFTPClientSession::parseAddress(const std::string& str, CSocketAddress& addr)
{
	std::string::const_iterator it  = str.begin();
	std::string::const_iterator end = str.end();
	while (it != end && *it != '(') ++it;
	if (it != end) ++it;
	std::string host;
	while (it != end && CAscii::isDigit(*it)) host += *it++;
	if (it != end && *it == ',') { host += '.'; ++it; }
	while (it != end && CAscii::isDigit(*it)) host += *it++;
	if (it != end && *it == ',') { host += '.'; ++it; }
	while (it != end && CAscii::isDigit(*it)) host += *it++;
	if (it != end && *it == ',') { host += '.'; ++it; }
	while (it != end && CAscii::isDigit(*it)) host += *it++;
	if (it != end && *it == ',') ++it;
	UInt16 portHi = 0;
	while (it != end && CAscii::isDigit(*it)) { portHi *= 10; portHi += *it++ - '0'; }
	if (it != end && *it == ',') ++it;
	UInt16 portLo = 0;
	while (it != end && CAscii::isDigit(*it)) { portLo *= 10; portLo += *it++ - '0'; }
	addr = CSocketAddress(host, portHi*256 + portLo);
}

void CFTPClientSession::parseExtAddress(const std::string& str, CSocketAddress& addr)
{
	std::string::const_iterator it  = str.begin();
	std::string::const_iterator end = str.end();
	while (it != end && *it != '(') ++it;
	if (it != end) ++it;
	char delim = '|';
	if (it != end) delim = *it++;
	if (it != end && *it == delim) ++it;
	if (it != end && *it == delim) ++it;
	UInt16 port = 0;
	while (it != end && CAscii::isDigit(*it)) { port *= 10; port += *it++ - '0'; }	
	addr = CSocketAddress(m_pControlSocket->peerAddress().host(), port);	
}

void CFTPClientSession::endTransfer()
{
	if (m_pDataStream)
	{
		delete m_pDataStream;
		m_pDataStream = 0;
		std::string response;
		int status = m_pControlSocket->receiveStatusMessage(response);
		if (!isPositiveCompletion(status)) 
			throw FTPException("Data transfer failed", response, status);
	}
}

NET_NAMESPACE_END
